import os

import cv2
import numpy as np

from option import parse_args
from tools.generate_utils import abs_error, flip, ssim, squared_error, get_real_fake_pairs
from tools.utils import hwc_to_chw, ssim_components

THRESHOLD_SUB = 0.1
THRESHOLD_SSIM = 0.25


def sub_gray2(diff_image_path, diff_image_path2):
    diff_image = cv2.imread(diff_image_path, cv2.IMREAD_GRAYSCALE)
    diff_image = diff_image / 255.
    diff_image[diff_image < THRESHOLD_SUB] = 0.
    diff_image[diff_image >= THRESHOLD_SUB] = 1.
    diff = (diff_image * 255).astype(np.uint8)
    cv2.imwrite(diff_image_path2, diff)


def ssim2(diff_image_path, diff_image_path2):
    diff_image = cv2.imread(diff_image_path, cv2.IMREAD_GRAYSCALE)
    diff_image = diff_image / 255.
    diff_image[diff_image < THRESHOLD_SSIM] = 0.
    diff_image[diff_image >= THRESHOLD_SSIM] = 1.
    diff = (diff_image * 255).astype(np.uint8)
    cv2.imwrite(diff_image_path2, diff)


def save_diff(args, mode='ssim'):
    real_fake_pairs = get_real_fake_pairs(args)
    for pair in real_fake_pairs:
        real_image, fake_image = pair
        s = fake_image.split('__')
        real_image_path = os.path.join(args.dst_dir, real_image)
        fake_image_path = os.path.join(args.dst_dir, fake_image)
        if os.path.exists(real_image_path) and os.path.exists(fake_image_path):
            img1 = cv2.imread(real_image_path, cv2.IMREAD_COLOR)
            img2 = cv2.imread(fake_image_path, cv2.IMREAD_COLOR)
            if mode == 'abs_error':
                diff_image = s[0] + '_abs_error_' + s[1]
                diff_image_path = os.path.join(args.dst_dir, diff_image)
                abs_error(img1, img2, diff_image_path)

                diff_image2 = s[0] + '_abs_errorz_' + s[1]
                diff_image_path2 = os.path.join(args.dst_dir, diff_image2)
                sub_gray2(diff_image_path, diff_image_path2)
            elif mode == 'squared_error':
                diff_image = f'{s[0]}_{s[1]}_{s[2]}' + '_squared_error_' + '{}_0.png'.format(int(s[3]))
                diff_image_path = os.path.join(args.dst_dir, diff_image)
                squared_error(img1, img2, diff_image_path)
            elif mode == 'flip':
                diff_image = fake_image[:-4] + 'flip_' + '{}_0.png'.format(int(fake_image.split('_')[2]))
                diff_image_path = os.path.join(args.dst_dir, diff_image)
                img1 = np.asarray(img1).astype(np.float32)
                img1 = hwc_to_chw(img1) / 255.
                img2 = np.asarray(img2).astype(np.float32)
                img2 = hwc_to_chw(img2) / 255.
                flip(img1, img2, diff_image_path)
            else:
                diff_image = s[0] + '_ssim_' + s[1]
                diff_image_path = os.path.join(args.dst_dir, diff_image)
                ssim(img1, img2, diff_image_path)

                diff_image_l = s[0] + '_ssiml_' + s[1]
                diff_image_l_path = os.path.join(args.dst_dir, diff_image_l)
                diff_image_c = s[0] + '_ssimc_' + s[1]
                diff_image_c_path = os.path.join(args.dst_dir, diff_image_c)
                diff_image_s = s[0] + '_ssims_' + s[1]
                diff_image_s_path = os.path.join(args.dst_dir, diff_image_s)
                ssim_components(img1, img2, diff_image_l_path, diff_image_c_path, diff_image_s_path)

                diff_image2 = s[0] + '_ssimz_' + s[1]
                diff_image_path2 = os.path.join(args.dst_dir, diff_image2)
                ssim2(diff_image_path, diff_image_path2)


def main():
    args = parse_args()
    args.dst_dir = os.path.join(args.dst_dir, 'celeb-df')
    # save_diff(args, mode='abs_error')
    # save_diff(args, mode='squared_error')
    save_diff(args, mode='ssim')
    # save_diff(args, mode='flip')


if __name__ == '__main__':
    main()
